"""
插件签名验证系统
提供插件数字签名和验证功能，防止插件被篡改
"""

import hashlib
import hmac
import json
import os
from typing import Dict, Optional
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives.serialization import load_pem_private_key, load_pem_public_key
from cryptography.exceptions import InvalidSignature


class PluginSignature:
    """插件签名信息"""
    
    def __init__(self, signature: bytes, public_key: bytes, timestamp: int):
        self.signature = signature
        self.public_key = public_key
        self.timestamp = timestamp


class SignatureManager:
    """签名管理器"""
    
    def __init__(self, keys_dir: str = "./keys"):
        self.keys_dir = keys_dir
        os.makedirs(keys_dir, exist_ok=True)
        self._loaded_keys: Dict[str, bytes] = {}
    
    def generate_key_pair(self, plugin_name: str) -> tuple[bytes, bytes]:
        """为插件生成密钥对"""
        # 生成RSA密钥对
        private_key = rsa.generate_private_key(
            public_exponent=65537,
            key_size=2048,
        )
        
        # 获取私钥PEM格式
        private_pem = private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption()
        )
        
        # 获取公钥PEM格式
        public_key = private_key.public_key()
        public_pem = public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        )
        
        # 保存密钥到文件
        private_key_path = os.path.join(self.keys_dir, f"{plugin_name}_private.pem")
        public_key_path = os.path.join(self.keys_dir, f"{plugin_name}_public.pem")
        
        with open(private_key_path, 'wb') as f:
            f.write(private_pem)
        
        with open(public_key_path, 'wb') as f:
            f.write(public_pem)
        
        return private_pem, public_pem
    
    def load_private_key(self, plugin_name: str) -> Optional[bytes]:
        """加载插件私钥"""
        key_path = os.path.join(self.keys_dir, f"{plugin_name}_private.pem")
        if os.path.exists(key_path):
            with open(key_path, 'rb') as f:
                return f.read()
        return None
    
    def load_public_key(self, plugin_name: str) -> Optional[bytes]:
        """加载插件公钥"""
        key_path = os.path.join(self.keys_dir, f"{plugin_name}_public.pem")
        if os.path.exists(key_path):
            with open(key_path, 'rb') as f:
                return f.read()
        return None
    
    def sign_plugin(self, plugin_name: str, plugin_data: bytes) -> PluginSignature:
        """为插件数据签名"""
        # 加载私钥
        private_pem = self.load_private_key(plugin_name)
        if not private_pem:
            raise ValueError(f"No private key found for plugin: {plugin_name}")
        
        private_key = load_pem_private_key(private_pem, password=None)
        
        # 生成签名
        signature = private_key.sign(
            plugin_data,
            padding.PSS(
                mgf=padding.MGF1(hashes.SHA256()),
                salt_length=padding.PSS.MAX_LENGTH
            ),
            hashes.SHA256()
        )
        
        # 获取公钥
        public_pem = self.load_public_key(plugin_name)
        if not public_pem:
            raise ValueError(f"No public key found for plugin: {plugin_name}")
        
        return PluginSignature(signature, public_pem, int(hashlib.time.time()))
    
    def verify_signature(self, plugin_data: bytes, signature: PluginSignature) -> bool:
        """验证插件签名"""
        try:
            # 加载公钥
            public_key = load_pem_public_key(signature.public_key)
            
            # 验证签名
            public_key.verify(
                signature.signature,
                plugin_data,
                padding.PSS(
                    mgf=padding.MGF1(hashes.SHA256()),
                    salt_length=padding.PSS.MAX_LENGTH
                ),
                hashes.SHA256()
            )
            
            return True
        except InvalidSignature:
            return False
        except Exception:
            return False
    
    def sign_plugin_files(self, plugin_name: str, plugin_dir: str) -> Dict[str, PluginSignature]:
        """为插件目录中的所有文件签名"""
        signatures = {}
        
        for root, dirs, files in os.walk(plugin_dir):
            for file in files:
                file_path = os.path.join(root, file)
                relative_path = os.path.relpath(file_path, plugin_dir)
                
                # 读取文件内容
                with open(file_path, 'rb') as f:
                    file_data = f.read()
                
                # 生成签名
                signature = self.sign_plugin(plugin_name, file_data)
                signatures[relative_path] = signature
        
        # 保存签名到签名文件
        signature_file = os.path.join(plugin_dir, "signature.json")
        signature_data = {}
        for path, sig in signatures.items():
            signature_data[path] = {
                "signature": sig.signature.hex(),
                "public_key": sig.public_key.decode('utf-8'),
                "timestamp": sig.timestamp
            }
        
        with open(signature_file, 'w') as f:
            json.dump(signature_data, f, indent=2)
        
        return signatures
    
    def verify_plugin_files(self, plugin_dir: str) -> bool:
        """验证插件目录中所有文件的签名"""
        signature_file = os.path.join(plugin_dir, "signature.json")
        if not os.path.exists(signature_file):
            return False
        
        # 读取签名文件
        with open(signature_file, 'r') as f:
            signature_data = json.load(f)
        
        # 验证每个文件
        for relative_path, sig_info in signature_data.items():
            file_path = os.path.join(plugin_dir, relative_path)
            if not os.path.exists(file_path):
                return False
            
            # 读取文件内容
            with open(file_path, 'rb') as f:
                file_data = f.read()
            
            # 构造签名对象
            signature = PluginSignature(
                bytes.fromhex(sig_info["signature"]),
                sig_info["public_key"].encode('utf-8'),
                sig_info["timestamp"]
            )
            
            # 验证签名
            if not self.verify_signature(file_data, signature):
                return False
        
        return True


class PluginSecurityManager:
    """插件安全管理器"""
    
    def __init__(self, keys_dir: str = "./keys"):
        self.signature_manager = SignatureManager(keys_dir)
        self._trusted_plugins: Dict[str, bytes] = {}
    
    def register_trusted_plugin(self, plugin_name: str, public_key: bytes):
        """注册受信任的插件"""
        self._trusted_plugins[plugin_name] = public_key
    
    def is_plugin_trusted(self, plugin_name: str) -> bool:
        """检查插件是否受信任"""
        return plugin_name in self._trusted_plugins
    
    def get_trusted_plugins(self) -> list[str]:
        """获取所有受信任的插件"""
        return list(self._trusted_plugins.keys())
    
    def verify_plugin_integrity(self, plugin_name: str, plugin_dir: str) -> bool:
        """验证插件完整性"""
        # 检查插件是否受信任
        if not self.is_plugin_trusted(plugin_name):
            return False
        
        # 验证文件签名
        return self.signature_manager.verify_plugin_files(plugin_dir)